# -*- coding:utf-8 -*-
import os,sys
import re
import traceback
import time
import pymongo
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), os.pardir))
from CommonLib.mylog import mylog
from CommonLib.db.MongoDBTool import MongoDBTool
from CommonLib.IndexCal.IndexCal import IndexCal
from CommonLib.IndexCal.FormatCal import FormatCal
import supeanut_config


'''
作者：supeanut
创建时间：2016-xx-xx xx:xx:xx
功能描述：
	xxx
	xxxxx
相关配置：
	supeanut_config.XXX
历史改动：
	2016-xx-xx: xxxxxx
'''
class StockSampling:
	def __init__(self):
		self.mongodb_obj = MongoDBTool()
		self.index_cal = IndexCal()
		self.for_cal = FormatCal()
		self.db = supeanut_config.MONGO_DB
		self.coll_day = supeanut_config.MONGO_COLL_DAYDATA
		self.coll_5min = supeanut_config.MONGO_COLL_5MINDATA
		self.all_data = None
		self.dimension_list = None
		self.target_list = None

	def caldims(self, period_s, period_e, adj, date_grain, code_str_list, dimension_list, target_list):
		self.dimension_list = dimension_list
		self.target_list = target_list
		dimension_str = 'dimensionGap'.join(dimension_list+target_list)
		dimension_info = self.for_cal.get_dimension_dict(dimension_str)
		code_dict = self.for_cal.cal_all_dims(dimension_info, period_s, period_e, code_str_list, adj, date_grain)
		inds = self.for_cal.get_inds(code_dict)
		self.all_data = inds

	def sampling(self, train_periods, valid_periods, test_periods):
		# 返回训练集，验证集，测试集
		# result_dict['train_set','valid_set','test_set']
		# train_set:{'features':[[dim1,dim2, ...], ..],labels:[[l1,l2,...],[],]
		field_name_dim = self.all_data['field_name_dim']
		features_train = []
		labels_train = []
		features_valid = []
		labels_valid = []
		features_test = []
		labels_test = []
		for ind in self.all_data['indivsiual']:
			datetime = ind[field_name_dim['datetime']]
			oneind_features = []
			oneind_labels = []
			if datetime >= train_periods[0] and train_periods[1] >= datetime:
				for feature_name in self.dimension_list:
					oneind_features.append(ind[field_name_dim[feature_name]])
				features_train.append(oneind_features)
				for label_name in self.target_list:
					oneind_labels.append(ind[field_name_dim[label_name]])
				labels_train.append(oneind_labels)
			if datetime >= valid_periods[0] and valid_periods[1] >= datetime:
				for feature_name in self.dimension_list:
					oneind_features.append(ind[field_name_dim[feature_name]])
				features_valid.append(oneind_features)
				for label_name in self.target_list:
					oneind_labels.append(ind[field_name_dim[label_name]])
				labels_valid.append(oneind_labels)
			if datetime >= test_periods[0] and test_periods[1] >= datetime:
				for feature_name in self.dimension_list:
					oneind_features.append(ind[field_name_dim[feature_name]])
				features_test.append(oneind_features)
				for label_name in self.target_list:
					oneind_labels.append(ind[field_name_dim[label_name]])
				labels_test.append(oneind_labels)
		result_dict = {'train_set':{},'valid_set':{},'test_set':{}}
		result_dict['train_set']['features'] = features_train
		result_dict['train_set']['labels'] = labels_train
		result_dict['valid_set']['features'] = features_valid
		result_dict['valid_set']['labels'] = labels_valid
		result_dict['test_set']['features'] = features_test
		result_dict['test_set']['labels'] = labels_test
		return result_dict

	# 包含当天，往前推pre_days个日期，（共pre_days个）
	def create_continue_index_name(self, index_name, params, pre_days):
		index_list = []
		for i in range(0,pre_days):
			param_str = ""
			for param in params:
				param_str += str(param) + ','
			if len(params) > 0:
				param_str = param_str[:-1]
				index_list.append(index_name + '(' + param_str + ',' + str(-i) + ')')
			else:
				index_list.append(index_name + '(' + str(-i) + ')')
		return index_list

if __name__ == '__main__':
	obj = StockSampling()
	ma_list = []
	for ma_num in [5, 10, 20, 55, 89, 89, 144, 233, 377, 450, 610]:
		ma_list += obj.create_continue_index_name('ma',[ma_num],20)
	print ma_list
	vol_list = []
	for vol_num in [5, 10, 20, 55, 89, 89, 144, 233, 377, 450, 610]:
		vol_list += obj.create_continue_index_name('volma_stock',[vol_num],20)
	print vol_list
	dimension_list = ma_list + ['open(0)','close(0)','high(0)','low(0)'] + vol_list
	target_list = ['(close(5)-close(0))/close(0)']
	obj.caldims(period_s = '2009-12-25 00:00:00', period_e = '2017-01-04 00:00:00', adj = 'ori', date_grain = 'day',\
				code_str_list = ['sz300033'],\
				dimension_list = dimension_list,\
				target_list = target_list,\
				)
	result_dict = obj.sampling(train_periods = ['2009-12-25 00:00:00', '2015-01-01 00:00:00'],\
				valid_periods=['2015-01-02 00:00:00', '2015-01-02 00:00:00'],\
				test_periods=['2015-01-02 00:00:00', '2017-01-04 00:00:00'])
	print '-------------------'
	print result_dict
